#! /usr/bin/bash


#1
#1 statBioFile 简单统计文件信息
#1 
#1 Usage:
#1     statBioFile [-h] [--help] file_path
#1 
#1 Options:
#1     -h,--help        查看帮助信息
#1 
#1 Example:
#1     $ statBioFile read.fq
#1


#2
#2 fq2fa 可以将 FASTQ 文件转换为 FASTA 文件，并提供一定的筛选能力。
#2       同时，也可以输入 FASTA 输出筛选后的 FASTA 。
#2
#2 Usage:
#2     fq2fa [-dltsh] [--help] input_file output_file
#2
#2 Options:
#2     -d,--delete_N    删除含有 N 碱基的序列
#2     -l,--less        删除长度小于给定 bp 的序列, default 90 
#2     -t,--terminal    删除两端给定 bp 的碱基, default 5 
#2     -s,--silent      安静模式，不显示运行提示 
#2     -h,--help        查看帮助信息
#2
#2 Example:
#2     $ fq2fa read.fq read.fa
#2     $ fq2fa -dl -t read.fa read_NON.fa
#2


#3
#3 revCmpDNAseq 可以对 DNA 序列进行反转互补，支持输入序列、FASTQ和FASTA文件。
#3              输入序列时可以直接输入，也可以从文件输入，默认不改变文件格式。
#3
#3 Usage:
#3     revCmpDNAseq [-rcosh] [--help] input_file output_file
#3
#3 Options:
#3     -r,--reverse         反转序列
#3     -c,--complementary   互补序列 
#3     -o,--output_format   输出的文件类型，可选seq、fasta、fastq     
#3                          选择 seq 是将得到仅仅碱基的文件，不包含信息行
#3                          注意无法将 seq 转为 fasta，同理无法将 fasta 转为 fastq
#3     -s,--silent          安静模式，不显示运行提示 
#3     -h,--help            查看帮助信息
#3
#3 Example:
#3     $ revCmpDNAseq -rc "ATCGATCG"
#3     $ revCmpDNAseq -rc -o "fasta" read.fq read.fa
#3


funclib_path=$(readlink -f ${BASH_SOURCE[0]})


function __checktype(){

    local filename="$1"
    
    if ! [[ -e "$filename" ]] ;then
        if [[ $filename =~ ^[ATCGNatcg]+$ ]] ;then
            return 10
        else
            echo "Can not find file: $filename"
            return 1
        fi
    fi

    local first_line=$(grep -v "^#" "$filename" | head -n 1)
    if egrep -q "^@" <<< "$first_line" ;then
        return 11
    elif egrep -q "^>" <<< "$first_line" ;then
        return 12
    elif egrep -q "^[ATCGNatcg]+$" <<< "$first_line" ;then
        return 15
    fi

    local columns23=$(echo "${first_line}" | cut -f2,3)
    if [[ "$columns23" =~ ^[a-z_A-Z]+[[:space:]][a-z_A-Z]+$ ]] ;then
        return 13
    elif [[ "$columns23" =~ ^[0-9]+[[:space:]][0-9]+$ ]] ;then
        return 14
    else
        return 20
    fi

}

function __statFile_head(){
    local size=$(du -h $1 |cut -f1)
    local filename=$(basename $1)
    printf "filename: %s\ttype: %s\tsize: %s\n" $filename $2 $size
}


function statBioFile(){

    if [[ "$1" == "-h" || "$1" == "--help" ]] ;then
        sed -rn 's/^#1 ?//;T;p;' "$funclib_path"
        return 0
    elif [[ "$1" =~ ^-.* ]] ;then
        echo "Usage: statBioFile [-h] [--help] file_path"
        return 1
    fi

    __checktype "$1"
    local type=$?
    if [[ $type -eq 11 || $type -eq 12 ]] ;then
        if [ $type -eq 11 ];then
            __statFile_head "$1" "FASTQ"
            local isfq=1
        else
            __statFile_head "$1" "FASTA"
            local isf1=0
        fi
        awk -v isfq=$isfq '
            BEGIN{
                a=c=t=g=n=0;
                sum_base=sum_seq=Nread_num=0;
                max_read_L=0;min_read_L=999999999
                max_GC=0;min_GC=1
            }\
            {
                if (NR%2==0 && ! (isfq && NR%4==0)){
                    sum_seq+=1
                    seqL=length($0)
                    sum_base+=seqL
                    if (seqL > max_read_L){
                        max_read_L=seqL
                    }
                    if (seqL < min_read_L ){
                        min_read_L=seqL
                    }
                    a+=gsub("A","");t+=gsub("T","")
                    seqC=gsub("C","");seqG=gsub("G","");
                    c+=seqC;g+=seqG
                    seqGC=(seqC+seqG)/seqL
                    if (seqGC > max_GC){
                        max_GC=seqGC
                    }
                    if (seqGC < min_GC){
                        min_GC=seqGC
                    }
                    seqN=gsub("N","")
                    n+=seqN
                    if (seqN > 0){
                        Nread_num++
                    }
                }
            } \
            END {
                printf("Reads\n ├─ sum: %s —— include N: %s\n",sum_seq,Nread_num)
                printf(\
                    " ├─ length avg: %0.2f —— min: %s —— max: %s \n",\
                    sum_base/sum_seq,min_read_L,max_read_L)
                printf(\
                    " └─ GC content —— min: %0.2f% —— max: %0.2f% \n",\
                    min_GC*100,max_GC*100)
                printf("Bases\n ├─ sum: %s —— N: %s\n",sum_base,n)
                printf(\
                    " ├─ A:%s %0.2f% —— T:%s %0.2f%\n └─ C:%s %0.2f% —— G:%s %0.2f%\n",\
                    a,a/sum_base*100,t,t/sum_base*100,\
                    c,c/sum_base*100,g,g/sum_base*100)
            }' "$1"
            return 0
    elif [[ $type -eq 13 || $type -eq 14 ]] ;then
        if [ $type -eq 13 ];then
            __statFile_head "$1" "GFF"
        else
            __statFile_head "$1" "BED"
        fi
        local row_num=$(grep -cv  "^[#[:space:]]" "$1")
        local col_num=$(grep -v "^#" "$1" | head -1 |tr "\t" "\n" | wc -l)
        echo "1.Shape: $row_num X $col_num" 
        printf "2.Chr -> count:\n"
        grep -v "^#" "$1" | cut -f1 | sort | uniq -c \
        | awk '{printf "  %s -- %s\n",$2,$1}' | column
        if [ $type -eq 13 ];then
            printf "3.Feature Type -> count:\n"
            grep -v "^#" "$1" | cut -f3 | sort | uniq -c \
            | awk '{printf "  %s -- %s\n",$2,$1}' | column
        fi
        return 0
    elif [[ $type -eq 20 || $type -eq 10 || $type -eq 15 ]] ;then
        echo "Unsupported file types."
        return 1
    else
        return 1
    fi
}

function fq2fa(){
    
    local rmN=false rmLowLine=false silent=false rmterminal=false

    tmp_parameters=$parameters
    parameters=$(getopt -o dshl::t:: \
                -l help,silent,delete_N,less::,terminal:: \
                -n "$FUNCNAME" -- "$@")
    [ $? -ne 0 ] && {
        echo "Usage: fq2fa [-dltsh] [--help] input_file output_file"
        unset parameters
        parameters=$tmp_parameters
        return 1
        }
    eval set -- "$parameters"   
    unset parameters
    parameters=$tmp_parameters

    while true ; do
        case "$1" in
            -h|--help)  
                sed -rn 's/^#2//;T;p;' "$funclib_path"
                return 0 ;;    
            -d|--delete_N) 
                rmN=true
                shift 1 ;;
            -s|--silent)
                silent=true
                shift 1 ;;
            -l|--less)
                rmLowLine=true    
                case "$2" in 
                    "")
                        local minlength=90
                        shift 2;;  
                    *)
                        local minlength="$2"
                        shift 2;;
                esac;;
            -t|--terminal)
                rmterminal=true
                case "$2" in 
                    "")
                        local rmtl=5
                        shift 2;;  
                    *)
                        local rmtl="$2"
                        shift 2;;  
                esac ;;
            --) 
                shift 1
                break ;;
            *) 
                echo "Parameters error, please check by --help"
                return 1 ;;
        esac
    done

    __checktype "$1"
    case $? in
        11) # The file is FASTQ
            awk '{if (NR%4==1) {sub(/^@/, ">"); print} \
                if (NR%4==2) {print}}' "$1" > "$2"
            ;;
        12) # The file in FASTA
            if $rmN || $rmLowLine || $rmterminal  ;then
                if ! $silent ;then  
                    echo "The file is FASTA and will be screened."
                fi
                cp "$1" "$2"
            else
                echo "The file is already FASTA."
                return 0
            fi
            ;;
        1)  
            return 1 
            ;;
        *)  
            echo "The file isn't FASTQ"
            return 1
            ;;
    esac

    if $rmN ;then
        cat "$2" | paste - - | awk -F '\t' '{if ($2 !~/N/) print}' \
            | sed "s/\t/\n/g" > fatmp
        mv fatmp "$2"
        if ! $silent ;then
            printf "Deletion of N-base-containing sequences completed.\n"
        fi
    fi

    if $rmLowLine ;then
        cat "$2" | paste - - | awk -F '\t' -v minL=$minlength '
            {if (length($2) >= minL) print}' | sed "s/\t/\n/g" > fatmp
        mv fatmp "$2"
        if ! $silent ;then
            printf "Deletion of sequences less than %s bp completed.\n" $minlength
        fi
    fi

    if $rmterminal ;then
        cat "$2" | awk -v rmtl=$rmtl '{
            if (NR%2==0) {seq=substr($0,rmtl+1,length($0)-rmtl*2);print seq} 
            else print}' > fatmp
        mv fatmp "$2"
        if ! $silent ;then
            printf "Excision of the terminal %s bp sequence completed.\n" $rmtl
        fi
    fi

    if ! $silent ;then
        echo "Finish! Please check file $(basename $2)"
    fi
    return 0
}

function revCmpDNAseq(){

    local reve=false comp=false silent=false supType=""

    tmp_parameters=$parameters
    parameters=$(getopt -o rcsho: \
                -l help,silent,reverse,comlementarys,output_format: \
                -n "$FUNCNAME" -- "$@")
    [ $? -ne 0 ] && { 
        echo "revCmpDNAseq [-rcosh] [--help] input_file output_file"
        unset parameters
        parameters=$tmp_parameters
        return 1
        }
    eval set -- "$parameters"   
    unset parameters
    parameters=$tmp_parameters
    
    while true ; do
        case "$1" in
            -h|--help)  
                sed -rn 's/^#3//;T;p;' "$funclib_path"
                return 0 ;;
            -s|--silent)
                silent=true
                shift 1 ;;
            -r|--reverse) 
                reve=true
                shift 1 ;;
            -c|--complementarys)
                comp=true
                shift 1 ;;
            -o|--output_format)
                case "$2" in 
                    "seq"|"fasta"|"fastq"|"")
                        local supType="$2"
                        shift 2 ;;
                    *)
                        printf "Only seq, fasta and fastq are supported.\n"
                        return 1;;
                esac ;;
            --)
                shift 1
                break ;;
            *) 
                echo "Parameters error, please check by --help"
                return 1 ;;
        esac
    done

    __checktype "$1"
    case $? in
        10) # Input is sequences
            if $reve && $comp ;then
                echo "$1" | rev | tr ATGCatgc TACGtacg
            elif $reve ;then
                echo "$1" | rev 
            elif $comp ;then
                echo "$1" | tr ATGCatgc TACGtacg
            else 
                echo "$1"
            fi
            return 0 ;;
        15) # The file is sequences
            case $supType in
                "seq"|"")
                    cp "$1" tmpfile_revCmp
                    supType="seq" ;;
                "fasta"|"fastq")
                    printf "Can not transfer to %s.\n" $supType
                    return 1 ;;
            esac ;;
        12) # The file is FASTA
            case $supType in
                "fasta"|"")
                    cp "$1" tmpfile_revCmp
                    supType="fasta" ;;
                "seq")
                    awk '{if (NR%2==0) print}' "$1" > tmpfile_revCmp ;;
                "fastq")
                    printf "Can not transfer to %s.\n" $supType
                    return 1 ;;
            esac ;;
        11) # The file is FASTQ
            case $supType in
                "fastq"|"")
                    cp "$1" tmpfile_revCmp
                    supType="fastq" ;;
                "fasta")
                    fq2fa -s "$1" tmpfile_revCmp ;;
                "seq")
                    awk '{if (NR%4==2) print}' "$1" > tmpfile_revCmp ;;
            esac ;;
        1)  return 1 ;;
        *) echo "This file type is not supported"
            return 1 ;;
    esac

    if [ "$supType" == "seq" ] ;then
        if $reve && $comp ;then
            rev tmpfile_revCmp | tr ATGCatgc TACGtacg > "$2"
        elif $reve ;then
            rev tmpfile_revCmp > "$2"
        elif $comp ;then
            cat tmpfile_revCmp | tr ATGCatgc TACGtacg > "$2"
        else 
            cat tmpfile_revCmp > "$2"
        fi 
    else
        local isfa=0 awk_reve=0 awk_comp=0
        if [ "$supType" == "fasta" ] ;then
            local isfa=1
        fi
        if $reve ;then local awk_reve=1 ;fi
        if $comp ;then local awk_comp=1 ;fi
        awk -v comp=$awk_comp \
            -v reve=$awk_reve \
            -v isfa=$isfa '
                function reverse(line){
                    revline=""
                    for (i=length(line);i>0;i--){
                        base=substr(line,i,1)
                        revline=revline base
                    }
                    return revline
                }
                function complementary(line){
                    seqforcomp="ACacgtGT"
                    for (i=1;i<5;i++){   
                        frombase=substr(seqforcomp,i,1)
                        tobase=substr(seqforcomp,9-i,1)
                        gsub(frombase,"P",line)
                        gsub(tobase,frombase,line)
                        gsub("P",tobase,line)
                    }
                    return line
                }
                {
                    if (NR%2==0) {
                        if (reve) {
                            $0=reverse($0)
                        }
                        if (isfa && comp) {
                            $0=complementary($0)
                        } else if (comp) {
                            if (NR%4==2) {
                                $0=complementary($0)
                            }
                        }
                        print 
                    } else {
                        print
                    }
                }' tmpfile_revCmp > "$2"
    fi
    
    rm tmpfile_revCmp
    if ! $silent ;then 
        echo "Finish! Please check file $(basename $2)"
    fi
    return 0
}

